from typing import List

import tqdm
import re
from langchain_community.document_loaders import UnstructuredFileLoader
from rag.module.indexing.loader.ocr import get_rapid_ocr
# from ocr import get_rapid_ocr

class CustomizedOcrPptLoader(UnstructuredFileLoader):
    def _get_elements(self) -> List:
        def ppt2text(filepath):
            from io import BytesIO

            import numpy as np
            from PIL import Image
            from pptx import Presentation
            from rapidocr_onnxruntime import RapidOCR

            # ocr = RapidOCR()
            ocr = get_rapid_ocr()
            prs = Presentation(filepath)
            resp = ""

            def extract_text(shape):
                nonlocal resp
                if shape.has_text_frame:
                    resp += shape.text.strip() + "\n"
                if shape.has_table:
                    for row in shape.table.rows:
                        for cell in row.cells:
                            for paragraph in cell.text_frame.paragraphs:
                                resp += paragraph.text.strip() + "\n"
                if shape.shape_type == 13:  # 13 表示图片
                    image = Image.open(BytesIO(shape.image.blob))
                    result, _ = ocr(np.array(image))
                    if result:
                        ocr_result = [line[1] for line in result]
                        resp += "\n".join(ocr_result)
                elif shape.shape_type == 6:  # 6 表示组合
                    for child_shape in shape.shapes:
                        extract_text(child_shape)

            b_unit = tqdm.tqdm(
                total=len(prs.slides), desc="RapidOCRPPTLoader slide index: 1"
            )
            # 遍历所有幻灯片
            for slide_number, slide in enumerate(prs.slides, start=1):
                b_unit.set_description(
                    "RapidOCRPPTLoader slide index: {}".format(slide_number)
                )
                b_unit.refresh()
                sorted_shapes = sorted(slide.shapes, key=lambda x: (x.top, x.left))  # 从上到下、从左到右遍历
                for shape in sorted_shapes:
                    extract_text(shape)
                b_unit.update(1)
            return resp

        text = ppt2text(self.file_path)
        text = text.replace("\n", "")  # 将换行符替换为空格
        # text = re.sub(r'\n+', '\n', text)
        from unstructured.partition.text import partition_text
        return partition_text(text=text, **self.unstructured_kwargs)


# # 创建一个 CustomizedOcrPptLoader 对象
# loader = CustomizedOcrPptLoader(file_path='/root/zzhou/data/开放原子teco-rag比赛数据/视频教学配套PPT/SDAA C编程模型/10 SDAA C编程模型-矩阵乘.pptx', unstructured_kwargs={})

# # 调用 _get_elements 方法获取元素
# elements = loader._get_elements()


# with open("/root/zzhou/ppt_txt.txt", 'w', encoding='utf-8') as output_file:
#     for element in elements:
#         # 将元素转换为字符串并写入文件
#         output_file.write(str(element) + '\n')